// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/exception/exception.h>
#include <pollux/dwio/dwrf/common/wrap/dwrf-proto-wrapper.h>

#include <boost/crc.hpp>
#define XXH_INLINE_ALL
#include <xxhash.h>

namespace kumo::pollux::dwrf {

class Checksum {
 public:
  explicit Checksum(proto::ChecksumAlgorithm type) : type_{type} {}

  virtual ~Checksum() = default;

  virtual void update(const void* input, size_t len) = 0;

  virtual int64_t getDigest(bool reset = true) = 0;

  proto::ChecksumAlgorithm getType() const {
    return type_;
  }

 private:
  proto::ChecksumAlgorithm type_;
};

class XxHash : public Checksum {
 public:
  XxHash()
      : Checksum{proto::ChecksumAlgorithm::XXHASH},
        state_{XXH64_createState()} {
    DWIO_ENSURE_NOT_NULL(state_, "failed to initialize xxhash");
    reset();
  }

  ~XxHash() override {
    XXH64_freeState(state_);
  }

  void update(const void* input, size_t len) override {
    DWIO_ENSURE_NE(XXH64_update(state_, input, len), XXH_ERROR);
  }

  int64_t getDigest(bool reset) override {
    auto ret = static_cast<int64_t>(XXH64_digest(state_));
    if (reset) {
      this->reset();
    }
    return ret;
  }

 private:
  XXH64_state_t* state_;
  unsigned long long seed_ = 0; // dwrf java uses seed 0

  void reset() {
    DWIO_ENSURE_NE(XXH64_reset(state_, seed_), XXH_ERROR);
  }
};

class Crc32 : public Checksum {
 public:
  Crc32() : Checksum{proto::ChecksumAlgorithm::CRC32} {
    reset();
  }

  ~Crc32() override = default;

  void update(const void* input, size_t len) override {
    crc_->process_bytes(input, len);
  }

  int64_t getDigest(bool reset) override {
    auto ret = crc_->checksum();
    if (reset) {
      this->reset();
    }
    return ret;
  }

 private:
  std::unique_ptr<boost::crc_32_type> crc_;

  void reset() {
    crc_ = std::make_unique<boost::crc_32_type>();
  }
};

class ChecksumFactory {
 public:
  static std::unique_ptr<Checksum> create(proto::ChecksumAlgorithm type) {
    switch (type) {
      case proto::ChecksumAlgorithm::NULL_:
        return {};
      case proto::ChecksumAlgorithm::XXHASH:
        return std::make_unique<XxHash>();
      case proto::ChecksumAlgorithm::CRC32:
        return std::make_unique<Crc32>();
      default:
        DWIO_RAISE("not supported");
    }
  }
};

} // namespace kumo::pollux::dwrf
